 
import os
import time
import asyncio
import aiohttp 
from typing import List,Tuple
 
from app.api.endpoints.doi import db_service
from app.services.downloader import DOIDownloader, JiangSuDownloader,NMPADownloader
from app.services.doi_parser import DOIParser
from app.core.config import settings

# from app.services.redis_service import redis_service
from app.services.db_service_factory import get_db_service
from app.models.database import NMPDDB,NMPDErrorDB, TaskDB, DOIItemDB

from playwright.async_api import BrowserContext ,async_playwright,Playwright,Browser
import re
db_service = get_db_service()


# 定义ARQ任务上下文的初始化和清理函数
async def startup(ctx):
    # 创建aiohttp会话，供所有任务共享
    ctx["session"] = aiohttp.ClientSession(
        headers={
            "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36 Edg/131.0.0.0",
            "Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8",
            "Accept-Language": "zh-CN,zh;q=0.9,en;q=0.8",
            "Accept-Encoding": "gzip, deflate, br",
            "Connection": "keep-alive",
            "Upgrade-Insecure-Requests": "1",
            "Sec-Fetch-Dest": "document",
            "Sec-Fetch-Mode": "navigate",
            "Sec-Fetch-Site": "none",
            "Sec-Fetch-User": "?1",
            "Cache-Control": "max-age=0",
            "sec-ch-ua": '"Not_A Brand";v="8", "Chromium";v="120", "Google Chrome";v="120"',
            "sec-ch-ua-mobile": "?0",
            "sec-ch-ua-platform": '"Windows"',
        }
    )


async def shutdown(ctx):
    # 关闭aiohttp会话
    if "session" in ctx:
        await ctx["session"].close()


# 检查任务是否被撤销
async def is_revoked(task_id: str):

    try:
        # task_data = redis_service.get(task_id)
        task_data = db_service.get_entity_as_dict(TaskDB, task_id)
        if not task_data:
            print(f"Redis中找不到任务 {task_id}，认为已被撤销")
            return True
    except Exception as e:
        print(f"is_revoked执行时出错: {str(e)}")
    return False


async def create_dir(task_id: str):

    # 从数据库获取任务信息，提取client_id和task_type
    task_info = db_service.get_entity_as_dict(TaskDB, task_id)
    client_id = task_info.get("client_id", "unknown")
    task_type = task_info.get("task_type", "DOI")  # 获取任务类型，默认为DOI

    # 确保client_id目录存在
    client_dir = os.path.join(settings.DOWNLOAD_DIR, client_id)
    os.makedirs(client_dir, exist_ok=True)

    # 创建任务类型子目录
    type_dir = os.path.join(client_dir, task_type)
    os.makedirs(type_dir, exist_ok=True)

    # 创建任务专用下载目录
    download_dir = os.path.join(type_dir, task_id)
    os.makedirs(download_dir, exist_ok=True)
    return client_id, client_dir, download_dir


# 测试任务
async def test_abc(ctx):
    print("nnnnnnnnnnnnnnnnnnnnnnnnnnnnnnnnnnnnnnnnnnn")
    return "测试成功"


# 主要的DOI处理任务
async def process_doi_list(ctx, task_id: str, doi_list: List[str]):
    import logging

    logger = logging.getLogger("uvicorn")

    try:
        logger.info(f"开始处理DOI列表任务，任务ID: {task_id}，共 {len(doi_list)} 个DOI")

        client_id, client_dir, download_dir = await create_dir(task_id)

        # 初始化任务状态
        task_status_update = {
            "status": "processing",
            "completed_dois": 0,
            "success_dois": 0,
            "failed_dois": 0,
        }

        # 更新任务状态
        await DOIParser.update_task_status(task_id, task_status_update)

        # 使用信号量控制并发
        semaphore = asyncio.Semaphore(settings.MAX_DOWNLOAD)
        counter_lock = asyncio.Lock()

        async def download_with_semaphore(doi: str):
            async with semaphore:
                try:
                    # 使用ctx中的session进行下载
                    success, file_path, error, metadata = (
                        await DOIDownloader.download_paper(
                            ctx["session"], task_id, doi, download_dir, is_revoked
                        )
                    )

                    async with counter_lock:
                        if success:
                            task_status_update["success_dois"] += 1
                        else:
                            await DOIParser.update_doi_metater(
                                task_id, doi, {"msg": error}
                            )
                            task_status_update["failed_dois"] += 1
                except asyncio.CancelledError:
                    pass
                except Exception as e:
                    async with counter_lock:
                        await DOIParser.update_doi_metater(
                            task_id, doi, {"msg": str(e)}
                        )
                        task_status_update["failed_dois"] += 1

                async with counter_lock:
                    task_status_update["completed_dois"] += 1
                    await DOIParser.update_task_status(task_id, task_status_update)

        # 创建并执行所有下载任务
        tasks = [download_with_semaphore(doi) for doi in doi_list]
        await asyncio.gather(*tasks)

        # 更新最终状态
        task_status_update["status"] = "completed"
        task_status_update["end_time"] = time.strftime("%Y-%m-%d %H:%M:%S")
        await DOIParser.update_task_status(task_id, task_status_update)

    except Exception as e:
        logger.error(f"处理DOI列表任务出错: {e}")
        task_status_update = {
            "status": "failed",
            "error": str(e),
            "end_time": time.strftime("%Y-%m-%d %H:%M:%S"),
        }
        await DOIParser.update_task_status(task_id, task_status_update)


async def process_jiangsu(ctx, task_id: str, url: str):
    try:
        # 初始化任务状态
        task_status_update = {
            "status": "processing",
            "completed_dois": 0,
            "success_dois": 0,
            "failed_dois": 0,
        }

        success, error, data = await JiangSuDownloader.download_table(
            ctx["session"], task_id, url
        )
        if success:
            client_id, client_dir, download_dir = await create_dir(task_id)
            items = data["items"]
            task_status_update["file_name"] = data["title"]
            task_status_update["total_dois"] = len(items)
            await DOIParser.update_task_status(task_id, task_status_update)
            # 使用信号量控制并发
            semaphore = asyncio.Semaphore(settings.MAX_DOWNLOAD)
            counter_lock = asyncio.Lock()

            async def download_with_semaphore(url: str, title: str, item_id: int):

                async with semaphore:
                    try:
                        # 使用ctx中的session进行下载
                        success, file_path, error, link, status = (
                            await JiangSuDownloader.dowload_file(
                                ctx["session"],
                                task_id,
                                url,
                                title,
                                download_dir,
                                is_revoked,
                            )
                        )

                        async with counter_lock:
                            if success:
                                await DOIParser.update_jiangsu_metater(
                                    task_id,
                                    item_id,
                                    {"status": "success", "pdf_url": link},
                                )
                                task_status_update["success_dois"] += 1
                            else:
                                await DOIParser.update_jiangsu_metater(
                                    task_id, item_id, {"msg": error}
                                )
                                task_status_update["failed_dois"] += 1
                    except asyncio.CancelledError:
                        pass
                    except Exception as e:
                        async with counter_lock:
                            await DOIParser.update_jiangsu_metater(
                                task_id, item_id, {"msg": str(e)}
                            )
                            task_status_update["failed_dois"] += 1

                    async with counter_lock:
                        task_status_update["completed_dois"] += 1
                        await DOIParser.update_task_status(task_id, task_status_update)

            # 创建并执行所有下载任务
            tasks = []
            metadata_list = []
            for index, item in enumerate(items):
                url = item["technical_report_link"]
                title = f'[{item["register_number"]}]{item["product_name"]}'
                metadata_list.append(
                    {
                        "task_id": task_id,
                        "title": title,
                        "status": "processing",
                        "doi": url,
                        "original_index": index,
                    }
                )

            entities = await DOIParser.add_jiangsu_metaer_list(
                task_id, metadata_list, False
            )
            await DOIParser.update_task_status(task_id, {"status": "processing"})

            for entity in entities:
                tasks.append(
                    download_with_semaphore(
                        entity["doi"], entity["title"], entity["id"]
                    )
                )

            await asyncio.gather(*tasks)
            # 更新最终状态

            task_status_update["status"] = "completed"
            task_status_update["end_time"] = time.strftime("%Y-%m-%d %H:%M:%S")
            await DOIParser.update_task_status(task_id, task_status_update)

        else:
            task_status_update["status"] = "failed"
            task_status_update["msg"] = error
            await DOIParser.update_task_status(task_id, task_status_update)
        # client_id,client_dir,download_dir=await create_dir(task_id)
    except Exception as e:
        task_status_update["msg"] = str(e)
        await DOIParser.update_task_status(task_id, task_status_update)


async def process_ylqx_bn(ctx, task_id: str, year: str,number:str,type:str,category:str,url:str,page_index:int=1,is_single_page:bool=False,is_single_rang=False):



    """
    处理医疗器械备案信息下载任务
    """
            # 初始化任务状态
    task_status_update = {
            "status": "processing",
            "completed_dois": 0,
            "success_dois": 0,
            "failed_dois": 0,
            "total_dois": 0,
            "year": year,
            "start_time": time.strftime("%Y-%m-%d %H:%M:%S"),
        }
    try:
        async with async_playwright() as p:
            browser,context=await create_playwright_context(p)  
            page = await context.new_page() 
            # 设置页面超时时间
            page.set_default_timeout(60000)  # 设置为60秒
            
            #将字符串映射成int
            category_map={
                "境内医疗器械（备案）":0,
                "境内医疗器械（备案历史数据）":1,
            }
            
            
            category_int= category_map[category]
            
            # newpage,total_count= await NMPADownloader.download_total_count(page,task_id,"https://www.nmpa.gov.cn/datasearch/home-index.html#category=ylqx", f'械备{year}',"境内医疗器械（备案）",is_revoked)
            
            success,flag,newpage,total_count= await NMPADownloader.download_total_count(page,task_id,url, f'{type}{year}',category,is_revoked) 
 
            number=str(number).zfill(3)
            is_loop=False
            is_error=False
            ttt=flag
            while total_count>0:
                    print(f"循环查询条件{number}")
                    try:
                      success2,flag2= await NMPADownloader.download_table_2(newpage,task_id, f'{type}{year}{number}',page_index, category_int, is_single_page,is_loop,total_count,is_revoked)
                      ttt=flag2
                      if success2 and flag2==0:
                          break
                       #newpage2,cu_table_count=await NMPADownloader.download_table(newpage,task_id, f'{type}{year}{number}',page_index, category_int, is_single_page,is_loop,total_count,is_revoked)
                    except Exception as e:
                        print(e) 
                        #判断key error_type是否在e对象里 
                        try:
                            error_data = e.args[0] if e.args and isinstance(e.args[0], dict) else {}
                            if error_data.get('error_type') == 'page_error':
                                print(e)
                                page_index = error_data.get('page_index')
                                is_error = True
                            
                        except (IndexError, AttributeError):
                            pass
                        finally:
                            number = str(int(number)-1).zfill(3)
                    if is_error is False: 
                        page_index=1         #大循环这里从第一页开始
                    is_loop=True
                    is_error=False
                    if is_single_page:
                            break
                    if int(number)<999 and is_single_rang is False:
                           number = str(int(number) + 1).zfill(3)
                           print(f"还要继续{total_count}:{number}:{page_index}")
                         
                    else:
                        break 
            print(f"搜索总结果数量: {total_count}")
           
            await page.close()
            # 关闭浏览器
            await browser.close()
            
            print(f"=============   {ttt}")
            
            if ttt!=0:
            # 更新最终状态
                task_status_update["status"] = "completed"
                task_status_update["completed_dois"] = total_count
                task_status_update["success_dois"] = total_count
                task_status_update["end_time"] = time.strftime("%Y-%m-%d %H:%M:%S")
                task_status_update["output_file"] = ''
                await DOIParser.update_task_status(task_id, task_status_update)

            return {
                "status": "success",
                "message": f"任务执行完成", 
            }

    except Exception as e:
        # 更新错误状态
        task_status_update = {
            "status": "failed",
            "error": str(e),
            "end_time": time.strftime("%Y-%m-%d %H:%M:%S"),
        }
        await DOIParser.update_task_status(task_id, task_status_update)

        return {"status": "error", "message": str(e)}
    
    
async def process_ylqx_bn_detail(ctx, task_id: str,url:str):
    try:
        async with async_playwright() as p:
            browser,context=await create_playwright_context(p)  
            await NMPADownloader.download_detail_tasks(context,url,task_id,is_revoked) 
            await browser.close()
         
    except Exception as e:
        print(e)
        return {"status": "error", "message": str(e)}

#返回组元
async def create_playwright_context(p:Playwright)->Tuple[Browser,BrowserContext]:   
        #async with async_playwright() as p:
            # 启动浏览器（尝试有头模式，便于调试）
            try_headless = False  # 设置为False可以在调试时查看浏览器界面

            # 添加浏览器启动参数 
            browser_args = [
                "--no-sandbox",
                "--disable-setuid-sandbox",
                "--disable-infobars",
                "--disable-dev-shm-usage",
                "--disable-blink-features=AutomationControlled",
                "--disable-gpu",
                "--disable-extensions",
                "--disable-client-side-phishing-detection",
                "--disable-translate",
                "--disable-notifications",
                "--disable-audio",
                "--disable-renderer-accessibility",
                "--disable-ipc-flooding-protection",
                "--enable-quic",
                "--disable-cache",
                "--ignore-certificate-errors",
                "--disable-logging",
                
            ]

            try:
                print(f"启动浏览器，headless={try_headless}")
                browser = await p.chromium.launch(
                    headless=try_headless, args=browser_args
                )
                # 创建更真实的浏览器上下文
                context = await browser.new_context( 
                    user_agent="Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/138.0.0.0 Safari/537.36 Edg/138.0.0.0", 
                    java_script_enabled=True,
                    ignore_https_errors=True, 
                    is_mobile=False,
                    accept_downloads=True,  
                )
                return browser,context
            except Exception as e:
                print(f"浏览器启动失败: {str(e)}")
                raise



    
    
    

